/*
 *  Copyright 2015 the original author or authors. 
 *  @https://github.com/scouter-project/scouter
 *
 *  Licensed under the Apache License, Version 2.0 (the "License"); 
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License. 
 */

package scouter.agent.asm.elasticsearch;

import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.LocalVariablesSorter;
import scouter.agent.ClassDesc;
import scouter.agent.Configure;
import scouter.agent.asm.IASM;
import scouter.agent.trace.TraceElasticSearch;

public class RestHighLevelClientASM implements IASM, Opcodes {

	private Configure conf = Configure.getInstance();

	public RestHighLevelClientASM() {
	}

	public ClassVisitor transform(ClassVisitor cv, String className, ClassDesc classDesc) {
		if (conf._hook_elasticsearch_enabled == false) {
			return cv;
		}
		if ("org/elasticsearch/client/RestHighLevelClient".equals(className) ||
				"org/opensearch/client/RestHighLevelClient".equals(className)) {

			return new RestHighLevelClientCV(cv, className);
		}
		return cv;
	}

	static class RestHighLevelClientCV extends ClassVisitor implements Opcodes {
		public String className;

		public RestHighLevelClientCV(ClassVisitor cv, String className) {
			super(ASM9, cv);
			this.className = className;
		}

		@Override
		public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
			MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
			if (mv == null) {
				return mv;
			}

			if ("internalPerformRequest".equals(name)) {
				return new RestHighLevelClientStartMV(access, desc, mv, className);
			}
			return mv;
		}
	}

	static class RestHighLevelClientStartMV extends LocalVariablesSorter implements Opcodes {
		private static final String TRACE = TraceElasticSearch.class.getName().replace('.', '/');
		private final static String METHOD = "startHighLevelRequest";
		private static final String SIGNATURE = "(Ljava/lang/Object;)Ljava/lang/Object;";
		private final static String END_METHOD = "endHighLevelRequest";
		private static final String END_SIGNATURE = "(Ljava/lang/Object;Ljava/lang/Throwable;)V";
		private Label startFinally = new Label();

		private String className;
		private int statIdx;

		public RestHighLevelClientStartMV(int access, String desc, MethodVisitor mv, String className) {
			super(ASM9, access, desc, mv);
			this.className = className;
		}

		@Override
		public void visitCode() {
			mv.visitVarInsn(ALOAD, 1);
			mv.visitMethodInsn(Opcodes.INVOKESTATIC, TRACE, METHOD, SIGNATURE, false);

			statIdx = newLocal(Type.getType(Object.class));
			mv.visitVarInsn(Opcodes.ASTORE, statIdx);
			mv.visitLabel(startFinally);

			mv.visitCode();
		}

		@Override
		public void visitInsn(int opcode) {
			if ((opcode >= IRETURN && opcode <= RETURN)) {
				mv.visitVarInsn(Opcodes.ALOAD, statIdx);
				mv.visitInsn(Opcodes.ACONST_NULL);
				mv.visitMethodInsn(Opcodes.INVOKESTATIC, TRACE, END_METHOD, END_SIGNATURE, false);
			}
			mv.visitInsn(opcode);
		}

		@Override
		public void visitMaxs(int maxStack, int maxLocals) {
			Label endFinally = new Label();
			mv.visitTryCatchBlock(startFinally, endFinally, endFinally, null);
			mv.visitLabel(endFinally);
			mv.visitInsn(DUP);
			int errIdx = newLocal(Type.getType(Throwable.class));
			mv.visitVarInsn(Opcodes.ASTORE, errIdx);

			mv.visitVarInsn(Opcodes.ALOAD, statIdx);
			mv.visitVarInsn(Opcodes.ALOAD, errIdx);
			mv.visitMethodInsn(Opcodes.INVOKESTATIC, TRACE, END_METHOD, END_SIGNATURE, false);
			mv.visitInsn(ATHROW);
			mv.visitMaxs(maxStack + 8, maxLocals + 2);
		}
	}

}
